{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Adding values to chain state\n",
    "\n",
    "The RunnablePassthrough.assign(...) static method takes an input value and adds the extra arguments passed to the assign function.<br>\n",
    "RunnablePassthrough.assign(...) static 方法采用输入值，并将传递给 assign 函数的额外参数相加。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'extra': {'num': 1, 'mult': 3}, 'modified': 2}"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_core.runnables import RunnableParallel, RunnablePassthrough\n",
    "\n",
    "runnable = RunnableParallel(\n",
    "    extra=RunnablePassthrough.assign(mult=lambda x: x[\"num\"] * 3),\n",
    "    modified=lambda x: x[\"num\"] + 1,\n",
    ")\n",
    "\n",
    "runnable.invoke({\"num\": 1})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Streaming\n",
    "One nice feature of this method is that it allows values to pass through as soon as they are available. To show this off, we'll use RunnablePassthrough.assign() to immediately return source docs in a retrieval chain:<br>\n",
    "此方法的一个很好的功能是，它允许值在可用时立即传递。为了展示这一点，我们将使用 RunnablePassthrough.assign() 在检索链中立即返回源文档："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'question': 'where did harrison work?'}\n",
      "{'context': [Document(page_content='harrison worked at kensho')]}\n",
      "{'output': ''}\n",
      "{'output': 'H'}\n",
      "{'output': 'arrison'}\n",
      "{'output': ' worked'}\n",
      "{'output': ' at'}\n",
      "{'output': ' Kens'}\n",
      "{'output': 'ho'}\n",
      "{'output': '.'}\n",
      "{'output': ''}\n",
      "{'output': ''}\n"
     ]
    }
   ],
   "source": [
    "from langchain_community.vectorstores import FAISS\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "from langchain_openai import ChatOpenAI, OpenAIEmbeddings\n",
    "import os\n",
    "\n",
    "os.environ[\"LANGCHAIN_PROJECT\"] = \"assign\"\n",
    "os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n",
    "os.environ[\"LANGCHAIN_API_KEY\"] = \"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\"\n",
    "\n",
    "vectorstore = FAISS.from_texts(\n",
    "    [\"harrison worked at kensho\"], embedding=OpenAIEmbeddings()\n",
    ")\n",
    "retriever = vectorstore.as_retriever()\n",
    "template = \"\"\"Answer the question based only on the following context:\n",
    "{context}\n",
    "\n",
    "Question: {question}\n",
    "\"\"\"\n",
    "prompt = ChatPromptTemplate.from_template(template)\n",
    "model = ChatOpenAI()\n",
    "\n",
    "generation_chain = prompt | model | StrOutputParser()\n",
    "\n",
    "retrieval_chain = {\n",
    "    \"context\": retriever,\n",
    "    \"question\": RunnablePassthrough(),\n",
    "} | RunnablePassthrough.assign(output=generation_chain)\n",
    "\n",
    "stream = retrieval_chain.stream(\"where did harrison work?\")\n",
    "\n",
    "for chunk in stream:\n",
    "    print(chunk)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchain0_1",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
